import numpy
import numpy as np
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import torch
import math
import prob


# 绘制三维矩形
def Rectangle_3D(ax, center, diff, facecolors = 'green', alpha = 1):
    '''
    :param ax: 绘图对象
    :param center: 方体中心 (x,y,z)
    :param diff: 2 * diff = (长 宽 高)
    diff 是一个元组 (长/2 , 宽/2, 高/2)
    :return: None
    '''

    x = center[0] - diff[0]
    y = center[1] - diff[1]
    z = center[2] - diff[2]

    vertices = np.array([[x, y, z],
                         [x + diff[0] * 2, y, z],
                         [x + diff[0] * 2, y + diff[1] * 2, z],
                         [x, y + diff[1] * 2, z],
                         [x, y, z + diff[2] * 2],
                         [x + diff[0] * 2, y, z + diff[2] * 2],
                         [x + diff[0] * 2, y + diff[1] * 2, z + diff[2] * 2],
                         [x, y + diff[1] * 2, z + diff[2] * 2]])

    # 定义长方体的六个面
    faces = [[vertices[0], vertices[1], vertices[2], vertices[3]],  # 底面
             [vertices[4], vertices[5], vertices[6], vertices[7]],  # 顶面
             [vertices[0], vertices[1], vertices[5], vertices[4]],  # 侧面
             [vertices[1], vertices[2], vertices[6], vertices[5]],  # 侧面
             [vertices[2], vertices[3], vertices[7], vertices[6]],  # 侧面
             [vertices[3], vertices[0], vertices[4], vertices[7]]]  # 侧面

    # 创建Poly3DCollection对象并添加到图形中
    poly3d = [faces[i] for i in range(len(faces))]
    ax.add_collection3d(Poly3DCollection(poly3d, facecolors=facecolors, linewidths=0.1, edgecolors='black', alpha=alpha))

# 绘制三维球体
def Circle_3D(ax, center, radius, facecolor = 'grey'):
    sita = np.linspace(0, 2 * np.pi, 20)
    fai = np.linspace(0, np.pi, 20)

    x = radius * np.outer(np.cos(sita), np.sin(fai)) + center[0]
    y = radius * np.outer(np.sin(sita), np.sin(fai)) + center[1]
    z = radius * np.outer(np.ones(np.size(sita)), np.cos(fai)) + center[2]

    ax.plot_surface(x, y, z, color=facecolor, alpha = 1)


def use_model(x,y,z):
    return (-0.54543733999999999 + math.sin(math.sin(((x - math.tanh((0.82548980000000005 + y))) / (1.1922542933451701 + 1.33552912280763 * math.exp(( - y + z)))))))

def plot_boundary_withoutModel(ax):
    # 生成一些采样的三维点和对应的函数值
    x = np.linspace(prob.domain_min[0], prob.domain_max[0], int(prob.PLOT_LEN_B[0]))
    # print("x:", x)  # x: [-3. -1.  1.  3.]
    y = np.linspace(prob.domain_min[1], prob.domain_max[1], int(prob.PLOT_LEN_B[1]))
    z = np.linspace(prob.domain_min[2], prob.domain_max[2], int(prob.PLOT_LEN_B[2]))

    # s = np.array(np.meshgrid(x, y, z))
    xx, yy, zz = np.meshgrid(x, y, z)
    s = np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T

    # print("s:",s)
    # print(type(s.size))
    print(s.size)
    #
    # nn_input = torch.tensor(s, dtype=torch.float64).to("cpu")
    nn_output = []
    for i in range((int)(s.size/3)):
        nn_output.append(use_model(s[i][0], s[i][1], s[i][2]))
    # nn_output = use_model(s[:, 0], s[:, 1], s[:, 2])
    # nn_output = use_model(xx, yy, zz)
    # plot_v = nn_output.detach().numpy()
    plot_v = numpy.array(nn_output)
    print("plot_v:", plot_v)
    # plot_v = np.squeeze(plot_v)

    # 找到函数值为零的点
    zero_indices = np.where(np.abs(plot_v) < 0.1)
    ss = s[zero_indices]
    # print(len(ss))
    size = np.ones(len(ss))
    # 绘制曲面
    ax.scatter(ss[:, 0], ss[:, 1], ss[:, 2], s=size, c='y', marker='o', alpha=0.2)


def plot_boundary(ax, model):
    # 生成一些采样的三维点和对应的函数值
    x = np.linspace(prob.domain_min[0], prob.domain_max[0], int(prob.PLOT_LEN_B[0]))
    y = np.linspace(prob.domain_min[1], prob.domain_max[1], int(prob.PLOT_LEN_B[1]))
    z = np.linspace(prob.domain_min[2], prob.domain_max[2], int(prob.PLOT_LEN_B[2]))

    # s = np.array(np.meshgrid(x, y, z))
    xx, yy, zz = np.meshgrid(x, y, z)
    s = np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T

    nn_input = torch.tensor(s, dtype=torch.float64).to("cpu")
    nn_output = model(nn_input)
    plot_v = nn_output.detach().numpy()
    plot_v = np.squeeze(plot_v)
    # print("plot_v:", plot_v)

    # 找到函数值为零的点
    zero_indices = np.where(np.abs(plot_v) < 0.1)
    ss = s[zero_indices]
    # print(len(ss))
    size = np.ones(len(ss))
    # 绘制曲面
    ax.scatter(ss[:, 0], ss[:, 1], ss[:, 2], s=size, c='y', marker='o', alpha=0.2)

    # ax.add_collection3d(Poly3DCollection([s[zero_indices]], facecolors='grey', edgecolors='yellow', alpha=0.1))

def plot_tran(ax, num, step):
    x_data = np.linspace(prob.init_min[0], prob.init_max[0], num)
    y_data = np.linspace(prob.init_min[1], prob.init_max[1], num)
    z_data = np.linspace(prob.init_min[2], prob.init_max[2], num)

    s = np.array(np.meshgrid(x_data, y_data, z_data))
    b = s.reshape(-1, order='F')
    s = b.reshape(-1, 3)
    x = s[:, 0]
    y = s[:, 1]
    z = s[:, 2]

    if prob.init_shape == 1:
        x0 = torch.tensor(s, dtype=torch.float32).to("cpu")
    if prob.init_shape == 2:
        p0 = (prob.init_min[0] + prob.init_max[0]) / 2.0
        p1 = (prob.init_min[1] + prob.init_max[1]) / 2.0
        p2 = (prob.init_min[2] + prob.init_max[2]) / 2.0
        r = p0 - prob.init_min[0]
        cons = (x - p0) * (x - p0) + (y - p1) * (y - p1) + (z - p2) * (z - p2) <= r * r
        x0 = torch.tensor(s[cons], dtype=torch.float32).to("cpu")
    print(x0.shape)
    # print(x0)

    x_pre = x0
    result = [x_pre]
    for i in range(step):
        x_next = prob.vector_field(x_pre)
        result.append(x_next)
        x_pre = x_next

    result = torch.cat(result, 1).cpu().detach().numpy()
    # print(result)

    for i in range(result.shape[0]):
        route = result[i, :].reshape((-1, 3))
        # print(route)
        a = route[:, 0]
        b = route[:, 1]
        c = route[:, 2]
        ax.scatter(a, b, c, alpha=1, s=1)

